﻿using Microsoft.ML.OnnxRuntime.Tensors;
using Microsoft.ML.OnnxRuntime;
using SixLabors.ImageSharp.Drawing.Processing;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using Newtonsoft.Json;
using Image = SixLabors.ImageSharp.Image;
using Font = SixLabors.Fonts.Font;
using SystemFonts = SixLabors.Fonts.SystemFonts;
using OpenCvSharp;
using Cv = OpenCvSharp.Cv2;

namespace OnnxNetFormTest
{
    internal class Model
    {
        static string modelPath = @"yolov5n.onnx";
        static int gpuDeviceId = 0;

        static InferenceSession session;
        static Image<Rgb24> image;

        static InferenceSession Session
        {
            get
            {
                if (session == null)
                {
                    session = new InferenceSession(modelPath, SessionOptions.MakeSessionOptionWithCudaProvider(gpuDeviceId));
                }
                return session;
            }
        }


        static Tensor<float> tensor = new DenseTensor<float>(new[] { 1, 3, 640, 640 });
        static Tensor<float> value;
        static Dictionary<int, List<Predict>> dict = new Dictionary<int, List<Predict>>();
        static IDisposableReadOnlyCollection<DisposableNamedOnnxValue>? results;
        static List<string> TypeNames = JsonConvert.DeserializeObject<List<string>>("['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']");

        public static void Calculate(string filepath)
        {
            image = Image.Load<Rgb24>(filepath);
            image.Mutate(x => x.Resize(640, 640));

            for (int j = 0; j < image.Height; j++)
            {
                image.ProcessPixelRows(img =>
                {
                    var pixelSpan = img.GetRowSpan(j);
                    for (int i = 0; i < image.Width; i++)
                    {
                        tensor[0, 0, j, i] = pixelSpan[i].R / 255f;
                        tensor[0, 1, j, i] = pixelSpan[i].G / 255f;
                        tensor[0, 2, j, i] = pixelSpan[i].B / 255f;
                    }
                });
            }

            // 将图片传至模型输入层
            // Setup inputs and outputs
            var inputs = new List<NamedOnnxValue>()
            {
                NamedOnnxValue.CreateFromTensor("images", tensor)
            };

            // 运行模型得到结果
            results = Session.Run(inputs);

            if (results.Any())
            {
                value = (DenseTensor<float>)results.ToArray().First().Value;
                dict.Clear();
                for (int i = 0; i < 25200; i++)
                {
                    //阈值
                    if (value[0, i, 4] >= 0.25)
                    {

                        int index = 5;
                        for (int j = 6; j < 85; j++)
                        {
                            if (value[0, i, j] > value[0, i, index])
                            {
                                index = j;
                            }
                        }
                        index -= 5;

                        Predict predict = new()
                        {
                            CenterX = value[0, i, 0],
                            CenterY = value[0, i, 1],
                            Width = value[0, i, 2],
                            Height = value[0, i, 3],
                            Confidence = value[0, i, 4],
                            Type = index
                        };

                        if (!dict.ContainsKey(index))
                        {
                            dict[index] = new List<Predict>();
                        }
                        dict[index].Add(predict);

                    }
                }

                foreach (int key in dict.Keys)
                {
                    List<Predict> predicts = dict[key];
                    predicts.Sort();

                    while (Predict.NMS(predicts, 0.5f)) ;

                    // 给检测的对象画框
                    // Put boxes, labels and confidence on image and save for viewing
                    using var outputImage = File.OpenWrite("image.jpg");
                    Font font = SystemFonts.CreateFont("Arial", 28);

                    foreach (Predict predict in predicts)
                    {
                        float centerX = predict.CenterX;
                        float centerY = predict.CenterY;
                        float width = predict.Width;
                        float height = predict.Height;

                        float x1 = centerX - width / 2,
                            y1 = centerY - height / 2,
                            x2 = centerX + width / 2,
                            y2 = centerY + height / 2;


                        image.Mutate(x =>
                        {
                            x.DrawLines(SixLabors.ImageSharp.Color.Red, 2f, new SixLabors.ImageSharp.PointF[]
                            {
                                new SixLabors.ImageSharp.PointF(x1, y1),
                                new SixLabors.ImageSharp.PointF(x1, y2),

                                new SixLabors.ImageSharp.PointF(x2, y2),
                                new SixLabors.ImageSharp.PointF(x2,y1),
                                new SixLabors.ImageSharp.PointF(x1, y1),
                            });
                            x.DrawText($"{TypeNames[predict.Type]}--{predict.Confidence:0.00}", font, SixLabors.ImageSharp.Color.Blue, new SixLabors.ImageSharp.PointF(x1, y1 + 3));
                        });

                        Console.WriteLine(JsonConvert.SerializeObject(predict));
                    }

                }

                // 图片保存到outputs文件夹下
                image.SaveAsJpeg("image.jpg");
            }

        }

        static Scalar scalar_rect = new Scalar(0, 0, 255);
        static Scalar scalar_font = new Scalar(255, 0, 0);
        //static Scalar scalar_type = new Scalar(0, 0, 255);

        public static void Calculate(Mat mat)
        {
            for (int i = 0; i < mat.Rows; i++)
            {
                for (int j = 0; j < mat.Cols; j++)
                {
                    Vec3b vec = mat.At<Vec3b>(i, j);
                    tensor[0, 0, i, j] = vec.Item2 / 255f;
                    tensor[0, 1, i, j] = vec.Item1 / 255f;
                    tensor[0, 2, i, j] = vec.Item0 / 255f;
                }
            }

            // 将图片传至模型输入层
            var inputs = new List<NamedOnnxValue>()
            {
                NamedOnnxValue.CreateFromTensor("images", tensor)
            };

            // 运行模型得到结果
            results = Session.Run(inputs);

            if (results.Any())
            {
                value = (DenseTensor<float>)results.ToArray().First().Value;
                dict.Clear();
                for (int i = 0; i < 25200; i++)
                {
                    //阈值
                    if (value[0, i, 4] >= 0.25)
                    {

                        int index = 5;
                        for (int j = 6; j < 85; j++)
                        {
                            if (value[0, i, j] > value[0, i, index])
                            {
                                index = j;
                            }
                        }
                        index -= 5;

                        Predict predict = new()
                        {
                            CenterX = value[0, i, 0],
                            CenterY = value[0, i, 1],
                            Width = value[0, i, 2],
                            Height = value[0, i, 3],
                            Confidence = value[0, i, 4],
                            Type = index
                        };

                        if (!dict.ContainsKey(index))
                        {
                            dict[index] = new List<Predict>();
                        }
                        dict[index].Add(predict);

                    }
                }

                foreach (int key in dict.Keys)
                {
                    List<Predict> predicts = dict[key];
                    predicts.Sort();

                    while (Predict.NMS(predicts, 0.5f)) ;

                    // 给检测的对象画框
                    Font font = SystemFonts.CreateFont("Arial", 28);

                    foreach (Predict predict in predicts)
                    {
                        float centerX = predict.CenterX;
                        float centerY = predict.CenterY;
                        float width = predict.Width;
                        float height = predict.Height;

                        float x = centerX - width / 2,
                            y = centerY - height / 2;

                        Cv.Rectangle(mat, new Rect((int)x, (int)y, (int)width, (int)height), scalar_rect, 1);
                        Cv.PutText(mat, $"{TypeNames[predict.Type]}--{predict.Confidence:0.00}", new OpenCvSharp.Point(x, y + 20), HersheyFonts.HersheyTriplex, 1, scalar_font);
                        //Cv.PutText(mat, $"{predict.Confidence}", new OpenCvSharp.Point(x, y + 50), HersheyFonts.HersheySimplex, 1, scalar_font);
                    }
                }
            }
        }

        private class Predict : IComparable<Predict>
        {
            public float CenterX { get; set; }
            public float CenterY { get; set; }
            public float Width { get; set; }
            public float Height { get; set; }

            public float Confidence { set; get; }
            public int Type { set; get; }

            public int CompareTo(Predict? other)
            {
                if (other?.Confidence - Confidence <= 0)
                {
                    return -1;
                }
                else
                {
                    return 1;
                }
            }

            public float IOU(Predict predict)
            {
                float w_2 = Width / 2;
                float h_2 = Height / 2;
                float x1 = CenterX - w_2, x2 = CenterX + w_2;
                float y1 = CenterY - h_2, y2 = CenterY + h_2;

                w_2 = predict.Width / 2;
                h_2 = predict.Height / 2;
                float x3 = predict.CenterX - w_2, x4 = predict.CenterX + w_2;
                float y3 = predict.CenterY - h_2, y4 = predict.CenterY + h_2;

                float minX = Math.Max(x1, x3);
                float maxX = Math.Min(x2, x4);
                float minY = Math.Max(y1, y3);
                float maxY = Math.Min(y2, y4);

                float w = Math.Max(0, maxX - minX);
                float h = Math.Max(0, maxY - minY);

                float jiao = w * h;
                float bing = Width * Height + predict.Width * predict.Height - jiao;
                return jiao / bing;
            }

            public static bool NMS(IList<Predict> predicts, float ioc_threshold = 0.5f)
            {
                bool flag = false;
                int index = -1;
                for (int i = 0; i < predicts.Count; i++)
                {

                    for (int j = i + 1; j < predicts.Count; j++)
                    {
                        float iou = predicts[i].IOU(predicts[j]);
                        if (iou >= ioc_threshold)
                        {
                            index = j;
                            flag = true;
                            break;
                        }
                    }
                    if (flag)
                    {
                        break;
                    }
                }
                if (flag)
                {
                    predicts.RemoveAt(index);
                }
                return flag;
            }
        }
    }
}
